import json

import numpy as np
from peewee import fn
from prefect import task
import pandas as pd
from prefect import get_run_logger
from functools import reduce
import operator
from application.constant import CN_LANGUAGE_TYPE, ARTICLE_TIME_PRECISION
from application.db.mysql_db.article_ch.CnArticleList import CnArticleList
from application.db.redis_db.base_redis import RedisManager
from application.tasks.base_task import BaseTask

redis_manager = RedisManager()
redis_conn = redis_manager.get_connection()


# BaseTask
class ArticleChInfoCleanTask(BaseTask):
    def run(self, df: pd.DataFrame):
        """
        清洗并计算文章信息的统计结果

        :param df: pandas.DataFrame
            输入的文献信息数据表，必须包含 `article_info` 字段
        :return: pandas.DataFrame
            返回清洗并增加检查字段后的 DataFrame
        """
        logger = get_run_logger()
        logger.info("Computing statistics")
        # 原地修改，节省内存
        # df.drop_duplicates(inplace=True)

        # 补充提取期刊编码
        df['article_info'] = df['article_info'].apply(self.safe_get_journal_code)

        # 检查文献信息是否完整
        df['check_article_info_does_it_exist'] = df['article_info'].apply(self.check_article_info)

        # 检查语言类型是否错误
        df['check_whether_the_language_type_is_wrong'] = df['article_info'].apply(self.is_language_type_wrong)

        # 检查中英文标题是否缺失
        df['check_chinese_and_english_titles_are_missing'] = df['article_info'].apply(self.are_titles_missing)

        # 检查发表时间精度是否异常
        df['check_published_time_precision'] = df['article_info'].apply(self.published_time_precision_wrong)

        # 检查是否存在期刊数据
        df['check_whether_the_journal_data_exists'] = df['article_info'].apply(self.check_journal_data)

        # 将中英文标题统一转换为 article_name 字段
        df['article_info'] = df['article_info'].apply(self.translate_article_name)

        # 提取用于判重的字段
        df['article_doi'] = df['article_info'].apply(lambda x: x.get('article_doi'))

        df['platform_id_and_platform_source'] = df['article_info'].apply(
            lambda x: {"platform_id": x.get('platform_id'), "platform_source": x.get('platform_source')}
        )

        df['article_chi_name'] = df['article_info'].apply(
            lambda x: {"article_chi_name": x.get('article_name', {}).get('chi'),
                       "article_eng_name": x.get('article_name', {}).get('eng'),
                       "journal_code": x.get('journal_info', {}).get("journal_code")}
            if x.get('article_name', {}).get('chi') else {}
        )
        df['article_eng_name'] = df['article_info'].apply(
            lambda x: {"article_eng_name": x.get('article_name', {}).get('eng'),
                       "journal_code": x.get('journal_info', {}).get("journal_code")}
        )
        # ========== 1) 通过 DOI 批量匹配 ==========
        df = self.doi_batch_matching(df)
        # mask = np.random.rand(len(df)) < 0.8
        # # 将随机一半的 article_id 置空
        # df.loc[mask, "article_id"] = np.nan
        # print(df['article_id'])
        # ========== 2) 通过 platform_id + platform_source 批量匹配 ==========
        # 只对仍未匹配到 article_id 的行进行尝试
        df = self.platform_batch_matching(df)
        # mask = np.random.rand(len(df)) < 0.8
        # # 将随机一半的 article_id 置空
        # df.loc[mask, "article_id"] = np.nan
        # print(df['article_id'])
        # ========== 3) 通过标题批量匹配（两个筛选：有中文标题的匹配；仅英文标题的匹配） ==========
        df = self.title_batch_matching(df)

        return df

    def doi_batch_matching(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        通过 DOI 批量匹配
        """
        query = CnArticleList.select(CnArticleList.article_id, CnArticleList.article_doi)
        # 过滤掉空 DOI
        doi_list = [d for d in df['article_doi'].tolist() if d]
        if doi_list:
            article_doi_query = query.where(CnArticleList.article_doi.in_(doi_list))
            doi_to_id = {item['article_doi']: item['article_id'] for item in article_doi_query.dicts()}
            df['article_id'] = df['article_doi'].map(doi_to_id)
        else:
            df['article_id'] = np.nan
        return df

    def platform_batch_matching(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        通过平台数据批量匹配
        """
        missing_platform_rows = df[df["article_id"].isna()]['platform_id_and_platform_source'].tolist()
        if missing_platform_rows:
            exprs = [
                (CnArticleList.platform_id == platform['platform_id']) &
                (CnArticleList.article_creator == platform['platform_source'])
                for platform in missing_platform_rows
                if platform.get('platform_id') is not None or platform.get('platform_source') is not None
            ]
            if exprs:
                platform_query = CnArticleList.select(CnArticleList.article_id,
                                                      CnArticleList.platform_id,
                                                      CnArticleList.article_creator).where(reduce(operator.or_, exprs))
                mapping = {
                    (item["platform_id"], item["article_creator"]): item["article_id"]
                    for item in platform_query.dicts()
                }
                df.loc[df['article_id'].isna(), 'article_id'] = df.loc[
                    df['article_id'].isna(), 'platform_id_and_platform_source'].apply(
                    lambda x: mapping.get((x.get('platform_id'), x.get('platform_source')))
                )
        return df

    def title_batch_matching(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        通过标题批量匹配
        :batch_title_matching:
        对仍未匹配到 article_id 的记录，分两类批量查询并回填：
        1) 有中文标题：以 (chi, eng) 为主要判重键，若有 journal_code 则同时作为限制；
        2) 仅英文标题：以 eng 为判重键，若有 journal_code 则同时作为限制。
        查询时使用 JSON_EXTRACT 将 article_name 中的 chi/eng 字段抽出来，以便在返回结果上构建映射。
        """
        # A. 有中文标题（中英文同时存在）
        need_title_match_idx = df['article_id'].isna()
        chi_mask = need_title_match_idx & df['article_chi_name'].apply(bool)
        chi_rows = df[chi_mask]

        if not chi_rows.empty:
            # 构造去重的条件集合（每个元素是一个 dict 包含 chi, eng, journal_code）
            unique_chi_tuples = chi_rows['article_chi_name'].apply(
                lambda x: (x.get('article_chi_name'), x.get('article_eng_name'), x.get('journal_code'))
            ).drop_duplicates().tolist()

            chi_exprs = []
            for chi, eng, jcode in unique_chi_tuples:
                # 跳过无效 title
                if not chi and not eng:
                    continue
                cond = (fn.JSON_EXTRACT(CnArticleList.article_name, '$.chi') == chi) & \
                       (fn.JSON_EXTRACT(CnArticleList.article_name, '$.eng') == eng)
                if jcode:
                    cond = cond & (CnArticleList.journal_code == jcode)
                chi_exprs.append(cond)

            if chi_exprs:
                chi_query = CnArticleList.select(
                    CnArticleList.article_id,
                    fn.JSON_EXTRACT(CnArticleList.article_name, '$.chi').alias('article_chi_name'),
                    fn.JSON_EXTRACT(CnArticleList.article_name, '$.eng').alias('article_eng_name'),
                    CnArticleList.journal_code
                ).where(reduce(operator.or_, chi_exprs))

                # 构建 (chi, eng, journal_code) -> article_id 的映射
                chi_mapping = {}
                for item in chi_query.dicts():
                    key = (item.get('article_chi_name'), item.get('article_eng_name'), item.get('journal_code'))
                    chi_mapping[key] = item.get('article_id')

                # 回填 df
                def map_chi_row(x):
                    if not x:
                        return None
                    key = (x.get('article_chi_name'), x.get('article_eng_name'), x.get('journal_code'))
                    return chi_mapping.get(key)

                df.loc[chi_mask, 'article_id'] = df.loc[chi_mask, 'article_chi_name'].apply(map_chi_row)

        # B. 仅英文标题（没有中文）
        eng_mask = need_title_match_idx & (~df['article_chi_name'].apply(bool))
        eng_rows = df[eng_mask]

        if not eng_rows.empty:
            unique_eng_tuples = eng_rows['article_eng_name'].apply(
                lambda x: (x.get('article_eng_name'), x.get('journal_code'))
            ).drop_duplicates().tolist()

            eng_exprs = []
            for eng, jcode in unique_eng_tuples:
                if not eng:
                    continue
                cond = (fn.JSON_EXTRACT(CnArticleList.article_name, '$.eng') == eng)
                if jcode:
                    cond = cond & (CnArticleList.journal_code == jcode)
                eng_exprs.append(cond)

            if eng_exprs:
                eng_query = CnArticleList.select(
                    CnArticleList.article_id,
                    fn.JSON_EXTRACT(CnArticleList.article_name, '$.eng').alias('article_eng_name'),
                    CnArticleList.journal_code
                ).where(reduce(operator.or_, eng_exprs))

                eng_mapping = {}
                for item in eng_query.dicts():
                    key = (item.get('article_eng_name'), item.get('journal_code'))
                    eng_mapping[key] = item.get('article_id')

                def map_eng_row(x):
                    if not x:
                        return None
                    key = (x.get('article_eng_name'), x.get('journal_code'))
                    return eng_mapping.get(key)

                df.loc[eng_mask, 'article_id'] = df.loc[eng_mask, 'article_eng_name'].apply(map_eng_row)
        return df

    def check_journal_data(self, article_info):
        article_info = article_info or {}
        journal_info = article_info.get("journal_info") or {}
        return bool(journal_info.get("journal_code"))

    def check_article_info(self, article_info):
        """
        判断文献信息是否完整

        :param article_info: dict
            单条文献信息数据
        :return: bool
            True 表示文献信息完整；False 表示缺失关键信息
        """
        article_info = article_info or {}
        if not article_info:
            return False

        journal_info = article_info.get("journal_info") or {}
        article_year = article_info.get("article_year")
        article_language = article_info.get("article_language")
        platform_source = article_info.get("platform_source")

        # “文献必填信息”检查条件：
        # 1. 期刊信息中至少包含 ISSN 或 期刊名称
        # 2. 文章年份存在
        # 3. 文章语言存在
        # 4. 数据来源平台存在
        return all([
            any([journal_info.get("journal_issn"), journal_info.get("journal_name")]),
            article_year,
            article_language,
            platform_source,
        ])

    def is_language_type_wrong(self, article_info):
        """
        判断语言类型是否错误

        :param article_info: dict
            单条文献信息数据
        :return: bool
            True 表示语言类型错误；False 表示正确或缺失
        """
        article_info = article_info or {}
        article_language = article_info.get("article_language")

        # 若未传语言，不视为“类型错误”
        if article_language is None:
            return False
        # 判断语言是否在允许集合内
        return article_language in CN_LANGUAGE_TYPE

    def are_titles_missing(self, article_info):
        """
        判断中英文标题是否同时缺失

        :param article_info: dict
            单条文献信息数据
        :return: bool
            True 表示中英文标题都不存在或为空；False 表示至少存在一个
        """
        article_info = article_info or {}
        article_en_name = article_info.get("article_en_name")
        article_cn_name = article_info.get("article_cn_name")
        return any([article_en_name, article_cn_name])

    def published_time_precision_wrong(self, article_info):
        """
        判断发表时间的精度值是否异常

        :param article_info: dict
            单条文献信息数据
        :return: bool
            True 表示精度异常；False 表示精度正确或允许缺失
        """
        article_info = article_info or {}
        published_time = article_info.get("published_time")
        published_time_precision = article_info.get("published_time_precision")

        # 若未传发表时间，视为异常
        if not published_time:
            return True

        # 若传入的精度值存在但不在允许集合中，则判为异常
        if published_time_precision:
            return published_time_precision in ARTICLE_TIME_PRECISION.values()

        # 若未传精度值，允许缺失
        return False

    def clc_codes_format_error(self, article_info):
        """
        判断中图分类号（clc_codes）的格式是否错误

        :param article_info: dict
            单条文献信息数据
        :return: bool
            True 表示 clc_codes 存在但不是 list 类型；False 表示格式正确或缺失
        """
        article_info = article_info or {}
        clc_codes = article_info.get("clc_codes")
        if clc_codes is None:
            return False
        return isinstance(clc_codes, list)

    def safe_get_journal_code(self, article_info):
        """
        安全提取期刊编码

        :param article_info: dict
            单条文献信息数据
        :return: str | None
            期刊编码；若缺失则返回 None
        """
        if not article_info:
            return None
        journal_info = article_info.get("journal_info") or {}
        journal_issn = journal_info.get("journal_issn")
        r_key = f"CN_JOURNAL_CODE_CACHE:{journal_issn}"
        journal_code = redis_conn.get(r_key)
        journal_info['journal_code'] = journal_code

        return article_info

    def translate_article_name(self, article_info):
        """
        将中英文标题转为统一的 article_name 字段

        :param article_info: dict
            单条文献信息数据
        :return: dict
            转换后的文献信息数据，新增 article_name 字段
        """
        article_en_name = article_info.pop("article_en_name", None)
        article_cn_name = article_info.pop("article_cn_name", None)
        article_name = {}

        # 如果存在英文标题，则添加至结构中
        if article_en_name:
            article_name['eng'] = article_en_name

        # 如果存在中文标题，则添加至结构中
        if article_cn_name:
            article_name['chi'] = article_cn_name

        article_info['article_name'] = article_name if article_name else {}
        return article_info
